"""
批量处理风机叶片图片 - Qwen版本
处理指定文件夹中的所有图片，生成YOLO标注和可视化图片
"""
import os
import time
from pathlib import Path
from collections import Counter
from typing import Dict, List
import shutil

import PIL.Image
from annotate_qwen import (
    CONFIG,
    preprocess_image,
    call_model,
    filter_defects,
    merge_nearby_defects,
    save_yolo_annotation,
    draw_boxes
)


# 批量处理配置
BATCH_CONFIG = {
    "input_dir": "风机叶片数据集/images",
    "output_dir": "风机叶片数据集/labels",
    "visualization_dir": "风机叶片数据集/visualizations",
    "skip_existing": False,  # 是否跳过已处理的图片
    "save_visualization": True,  # 是否保存可视化图片
    "delay_between_requests": 3,  # API请求间隔（秒），避免速率限制
    "save_summary": True,  # 保存统计报告
}


class QwenBatchProcessor:
    """批量处理器"""
    
    def __init__(self):
        self.input_dir = Path(BATCH_CONFIG["input_dir"])
        self.output_dir = Path(BATCH_CONFIG["output_dir"])
        self.viz_dir = Path(BATCH_CONFIG["visualization_dir"])
        
        # 创建输出目录
        self.output_dir.mkdir(parents=True, exist_ok=True)
        if BATCH_CONFIG["save_visualization"]:
            self.viz_dir.mkdir(parents=True, exist_ok=True)
        
        # 统计信息
        self.stats = {
            "total": 0,
            "processed": 0,
            "skipped": 0,
            "failed": 0,
            "no_defects": 0,
            "defect_counts": Counter()
        }
        
        # API密钥
        self.api_key = os.environ.get("DASHSCOPE_API_KEY")
        if not self.api_key:
            raise ValueError("未找到DASHSCOPE_API_KEY环境变量")
    
    def get_image_files(self) -> List[Path]:
        """获取所有图片文件"""
        extensions = ['.jpg', '.jpeg', '.png', '.bmp']
        files = []
        for ext in extensions:
            files.extend(self.input_dir.glob(f"*{ext}"))
            files.extend(self.input_dir.glob(f"*{ext.upper()}"))
        return sorted(files)
    
    def should_skip(self, image_path: Path) -> bool:
        """判断是否应该跳过此图片"""
        if not BATCH_CONFIG["skip_existing"]:
            return False
        
        label_path = self.output_dir / f"{image_path.stem}.txt"
        return label_path.exists()
    
    def process_single_image(self, image_path: Path) -> bool:
        """处理单张图片，返回是否成功"""
        try:
            # 1. 读取图像
            image = PIL.Image.open(image_path)
            print(f"   尺寸: {image.width}x{image.height}")
            
            # 2. 预处理
            print(f"   [1/5] 图像预处理...")
            processed = preprocess_image(image)
            
            # 3. AI分析
            print(f"   [2/5] AI模型分析...")
            defects = call_model(processed, self.api_key)
            
            if defects is None:
                print(f"   ❌ API调用失败")
                return False
            
            print(f"   [3/5] 原始结果: {len(defects)} 个")
            
            # 4. 过滤
            defects = filter_defects(defects)
            print(f"   [4/5] 过滤后: {len(defects)} 个")
            
            # 5. 合并
            defects = merge_nearby_defects(defects)
            print(f"   [5/5] 合并后: {len(defects)} 个")
            
            # 6. 保存结果
            if not defects:
                print(f"   ℹ️  未检测到有效缺陷")
                # 创建空标注文件
                label_path = self.output_dir / f"{image_path.stem}.txt"
                label_path.write_text("", encoding="utf-8")
                print(f"   💾 已保存空标注文件")
                
                # 保存原图副本作为可视化
                if BATCH_CONFIG["save_visualization"]:
                    viz_path = self.viz_dir / f"{image_path.stem}_annotated.jpg"
                    try:
                        shutil.copy(str(image_path), str(viz_path))
                        print(f"   ✅ 已保存原图副本: {viz_path.name}")
                    except Exception as e:
                        print(f"   ❌ 复制原图失败: {e}")
                
                self.stats["no_defects"] += 1
                self.stats["processed"] += 1
                return True
            
            # 统计缺陷类型
            counts = Counter(d["defect_type"] for d in defects)
            self.stats["defect_counts"].update(counts)
            
            print(f"   ✅ 检测完成:")
            for dtype, count in sorted(counts.items()):
                print(f"      - {dtype}: {count} 个")
            
            # 7. 保存YOLO标注
            label_path = self.output_dir / f"{image_path.stem}.txt"
            save_yolo_annotation(defects, label_path)
            print(f"   💾 已保存标注文件")
            
            # 8. 保存可视化
            if BATCH_CONFIG["save_visualization"]:
                viz_path = self.viz_dir / f"{image_path.stem}_annotated.jpg"
                draw_boxes(image_path, defects, viz_path)
            
            self.stats["processed"] += 1
            return True
            
        except Exception as e:
            print(f"   ❌ 处理失败: {e}")
            self.stats["failed"] += 1
            return False
    
    def run(self):
        """执行批量处理"""
        print("=" * 70)
        print("风力发电叶片缺陷批量检测系统 (Qwen版)")
        print("=" * 70)
        print()
        
        # 获取图片列表
        image_files = self.get_image_files()
        self.stats["total"] = len(image_files)
        
        if not image_files:
            print(f"❌ 在 {self.input_dir} 中未找到图片文件")
            return
        
        print(f"📂 输入目录: {self.input_dir}")
        print(f"📂 标注输出: {self.output_dir}")
        print(f"📂 可视化输出: {self.viz_dir}")
        print(f"📊 找到图片: {len(image_files)} 张")
        print(f"⏱️  请求间隔: {BATCH_CONFIG['delay_between_requests']} 秒")
        print()
        print("=" * 70)
        print()
        
        # 处理每张图片
        for idx, image_path in enumerate(image_files, 1):
            print("=" * 70)
            print(f"📷 [{idx}/{len(image_files)}] {image_path.name}")
            print(f"   进度: {idx/len(image_files)*100:.1f}%")
            print()
            
            # 检查是否跳过
            if self.should_skip(image_path):
                print(f"   ⏭️  已存在，跳过")
                self.stats["skipped"] += 1
                continue
            
            # 处理图片
            success = self.process_single_image(image_path)
            
            # API调用间隔（除了最后一张）
            if idx < len(image_files) and success:
                delay = BATCH_CONFIG["delay_between_requests"]
                if delay > 0:
                    print(f"   ⏳ 等待 {delay} 秒...")
                    time.sleep(delay)
            
            print()
        
        # 打印统计报告
        self.print_summary()
        
        # 保存统计报告
        if BATCH_CONFIG["save_summary"]:
            self.save_summary()
    
    def print_summary(self):
        """打印统计报告"""
        print()
        print("=" * 70)
        print("📊 处理统计报告")
        print("=" * 70)
        print(f"总图片数:     {self.stats['total']} 张")
        print(f"成功处理:     {self.stats['processed']} 张")
        print(f"跳过:         {self.stats['skipped']} 张")
        print(f"失败:         {self.stats['failed']} 张")
        print(f"未检测到缺陷: {self.stats['no_defects']} 张")
        print()
        
        if self.stats["defect_counts"]:
            print("缺陷类型统计:")
            for dtype, count in sorted(self.stats["defect_counts"].items(), 
                                      key=lambda x: x[1], reverse=True):
                print(f"  - {dtype}: {count} 个")
        else:
            print("未检测到任何缺陷")
        
        print("=" * 70)
    
    def save_summary(self):
        """保存统计报告到文件"""
        summary_path = Path("batch_summary_qwen.txt")
        
        lines = [
            "=" * 70,
            "风力发电叶片缺陷批量检测统计报告 (Qwen版)",
            "=" * 70,
            "",
            f"处理时间: {time.strftime('%Y-%m-%d %H:%M:%S')}",
            f"输入目录: {self.input_dir}",
            f"标注输出: {self.output_dir}",
            f"可视化输出: {self.viz_dir}",
            "",
            "=" * 70,
            "处理统计",
            "=" * 70,
            f"总图片数:     {self.stats['total']} 张",
            f"成功处理:     {self.stats['processed']} 张",
            f"跳过:         {self.stats['skipped']} 张",
            f"失败:         {self.stats['failed']} 张",
            f"未检测到缺陷: {self.stats['no_defects']} 张",
            "",
        ]
        
        if self.stats["defect_counts"]:
            lines.extend([
                "=" * 70,
                "缺陷类型统计",
                "=" * 70,
            ])
            for dtype, count in sorted(self.stats["defect_counts"].items(), 
                                      key=lambda x: x[1], reverse=True):
                lines.append(f"{dtype}: {count} 个")
        
        lines.append("=" * 70)
        
        summary_path.write_text("\n".join(lines), encoding="utf-8")
        print(f"\n📄 统计报告已保存: {summary_path}")


def main():
    """主函数"""
    try:
        processor = QwenBatchProcessor()
        processor.run()
    except ValueError as e:
        print(f"❌ 错误: {e}")
        print("\n请先设置API密钥:")
        print('$env:DASHSCOPE_API_KEY = "your-api-key"')
    except KeyboardInterrupt:
        print("\n\n⚠️  用户中断，正在安全退出...")
    except Exception as e:
        print(f"\n❌ 发生错误: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()

